{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c2e72e87",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 831, 10, 10])\n",
      "torch.Size([16, 831, 10, 10])\n",
      "编码值: 0, 原标签: Air\n",
      "编码值: 1, 原标签: Ea\n",
      "编码值: 2, 原标签: H2\n",
      "torch.Size([53184, 8, 10])\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import mean_absolute_error, mean_squared_error\n",
    "import random\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "# print(f\"Using device: {device}\")\n",
    "# model = torch.load('my_multitask_model.pth', map_location=device)\n",
    "\n",
    "np.set_printoptions(threshold=np.inf)\n",
    "torch.set_printoptions(threshold=torch.inf)\n",
    "\n",
    "validation_size = 0.2\n",
    "\n",
    "def set_seed(seed):\n",
    "    # seed init.\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "\n",
    "    # torch seed init.\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "#     torch.backends.cudnn.enabled = False \n",
    "#     os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':16:8'\n",
    "#     torch.use_deterministic_algorithms(True)\n",
    "    \n",
    "set_seed(42)\n",
    "def data_process(data, window_size, step):\n",
    "    tensor_data = torch.from_numpy(data)\n",
    "    window_split = tensor_data.unfold(1, window_size, step)  # parameter：dimension = 1\n",
    "#     (batch size, time step, feature&label) to (batch size, num_windows, feature&label, window_size)\n",
    "    window_data = window_split.numpy()\n",
    "\n",
    "#     window_data_list = window_data.tolist()\n",
    "#     random.shuffle(window_data_list)\n",
    "#     window_data = np.array(window_data_list)   \n",
    "    \n",
    "    # 划分 Feature 和 Label\n",
    "    print(window_split.shape)\n",
    "    features = window_data[:, :, :8, :]  # (64, num_windows, 8, window_size)\n",
    "    classification_labels = window_data[:, :, 8, -1]  # 二维\n",
    "    regression_labels = window_data[:, :, 9, -1]  # 二维\n",
    "    \n",
    "    # 转换为 tensor\n",
    "    feature_tensor = torch.tensor(features, dtype=torch.float32)\n",
    "    classification_label_tensor = torch.tensor(classification_labels, dtype=torch.long)\n",
    "    regression_label_tensor = torch.tensor(regression_labels, dtype=torch.float32)\n",
    "    \n",
    "    return feature_tensor, classification_label_tensor, regression_label_tensor\n",
    "\n",
    "# 样本数据\n",
    "X = np.zeros((80, 840, 10))  # 包括特征和标签\n",
    "sample_counts = {\n",
    "    '100ppmH2': 5, '200ppmH2': 5, '300ppmH2': 5, '400ppmH2': 5,\n",
    "    '500ppmH2': 5, '600ppmH2': 5, '700ppmH2': 5, '800ppmH2': 5,\n",
    "    '900ppmH2': 5, '1000ppmH2': 5, '90ppmEa': 5, '100ppmEa': 5,\n",
    "    '200ppmEa': 5, '300ppmEa': 5, '400ppmEa': 5, '500ppmEa': 5\n",
    "}\n",
    "sample_index = 0   \n",
    "label_encoder = LabelEncoder()\n",
    "\n",
    "# 收集所有分类标签进行编码\n",
    "all_labels = []\n",
    "for sample_type, count in sample_counts.items():\n",
    "    for j in range(count):\n",
    "        filename = rf'C:\\Users\\52898\\Desktop\\2024\\competition\\sampleV1.1\\{sample_type}_{j+1}.csv'\n",
    "        data = pd.read_csv(filename)\n",
    "        all_labels.extend(data.iloc[:840, 9])  # 收集第10列的分类标签\n",
    "label_encoder.fit(all_labels)\n",
    "\n",
    "for sample_type, count in sample_counts.items():\n",
    "    for j in range(count):\n",
    "        filename = rf'C:\\Users\\52898\\Desktop\\2024\\competition\\sampleV1.1\\{sample_type}_{j+1}.csv'\n",
    "        data = pd.read_csv(filename)\n",
    "        # 取第1到第11列，包括特征和标签\n",
    "        features = data.iloc[:, 1:11].values  \n",
    "        features = features[:840, :]  # 截取第1到第840行\n",
    "        # 编码分类标签（第10列）\n",
    "        features[:, 8] = label_encoder.transform(features[:, 8])\n",
    "        X[sample_index, :, :] = features\n",
    "        sample_index += 1\n",
    "\n",
    "# 划分特征和标签\n",
    "features = X[:, :, :8]  # 第1到第8列作为特征\n",
    "labels = X[:, :, 8:10]  # 第9和第10列作为标签\n",
    "\n",
    "# 划分训练集和验证集\n",
    "X_train, X_validation, labels_train, labels_validation = train_test_split(features, labels, shuffle=True, test_size=0.2, random_state=42)\n",
    "\n",
    "Y_classification_train = labels_train[:, :, 0]\n",
    "Y_classification_validation = labels_validation[:, :, 0]\n",
    "Y_regression_train = labels_train[:, :, 1]\n",
    "Y_regression_validation = labels_validation[:, :, 1]\n",
    "\n",
    "# 标准化\n",
    "X_scaler = StandardScaler()\n",
    "y_scaler = StandardScaler()\n",
    "\n",
    "X_train = X_train.reshape(-1, 8)\n",
    "X_train = X_scaler.fit_transform(X_train)\n",
    "X_train = X_train.reshape(-1, 840, 8)\n",
    "\n",
    "X_validation = X_validation.reshape(-1, 8)\n",
    "X_validation = X_scaler.transform(X_validation)\n",
    "X_validation = X_validation.reshape(-1, 840, 8)\n",
    "\n",
    "Y_regression_train = y_scaler.fit_transform(Y_regression_train.reshape(-1, 1)).reshape(-1, 840)\n",
    "Y_regression_validation = y_scaler.transform(Y_regression_validation.reshape(-1, 1)).reshape(-1, 840)\n",
    "\n",
    "# 滑动窗口参数\n",
    "window_size = 10 # 窗口大小\n",
    "step = 1  # 步长\n",
    "\n",
    "# 执行滑动窗口处理\n",
    "train_data = np.concatenate((X_train, Y_classification_train[:, :, np.newaxis], Y_regression_train[:, :, np.newaxis]), axis=2)\n",
    "train_features, train_classification_labels, train_regression_labels = data_process(train_data, window_size, step)\n",
    "\n",
    "validation_data = np.concatenate((X_validation, Y_classification_validation[:, :, np.newaxis], Y_regression_validation[:, :, np.newaxis]), axis=2)\n",
    "validation_features, validation_classification_labels, validation_regression_labels = data_process(validation_data, window_size, step)\n",
    "\n",
    "for index, label in enumerate(label_encoder.classes_):\n",
    "    print(f\"编码值: {index}, 原标签: {label}\")\n",
    "\n",
    "train_features = train_features.reshape(-1, 8, 10) #input shape = (batch size, feature, time step)\n",
    "train_classification_labels = train_classification_labels.reshape(-1)\n",
    "train_regression_labels = train_regression_labels.reshape(-1,1)\n",
    "validation_features = validation_features.reshape(-1, 8, 10)\n",
    "validation_classification_labels = validation_classification_labels.reshape(-1)\n",
    "validation_regression_labels = validation_regression_labels.reshape(-1,1)\n",
    "print(train_features.shape)\n",
    "# original_features = X_scaler.inverse_transform(train_features.permute(0, 1, 3, 2).reshape(-1, 8))\n",
    "\n",
    "# 创建数据集和数据加载器\n",
    "train_dataset = TensorDataset(train_features, train_regression_labels, train_classification_labels)\n",
    "validation_dataset = TensorDataset(validation_features, validation_regression_labels, validation_classification_labels)\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
    "validation_loader = DataLoader(validation_dataset, batch_size=32, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "aecdeb9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Scaler parameters saved to scalers.json\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "\n",
    "# 提取 X_scaler 的缩放参数\n",
    "X_scaler_params = {\n",
    "    \"mean_\": X_scaler.mean_.tolist(),\n",
    "    \"scale_\": X_scaler.scale_.tolist()\n",
    "}\n",
    "\n",
    "# 提取 y_scaler 的缩放参数\n",
    "y_scaler_params = {\n",
    "    \"mean_\": y_scaler.mean_.tolist(),\n",
    "    \"scale_\": y_scaler.scale_.tolist()\n",
    "}\n",
    "\n",
    "# 将两个 scaler 参数存储在一个字典中\n",
    "scaler_params = {\n",
    "    \"X_scaler\": X_scaler_params,\n",
    "    \"y_scaler\": y_scaler_params\n",
    "}\n",
    "\n",
    "# 将字典保存为 JSON 文件\n",
    "with open('scalers.json', 'w') as f:\n",
    "    json.dump(scaler_params, f, indent=4)\n",
    "\n",
    "print(\"Scaler parameters saved to scalers.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f00ead71",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyMultiTaskModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MyMultiTaskModel, self).__init__()\n",
    "        self.conv1 = nn.Conv1d(in_channels=8, out_channels=8, kernel_size=3, stride=1, padding=1)\n",
    "        self.pool = nn.MaxPool1d(kernel_size=2)\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.fc1 = nn.Linear(40, 10)\n",
    "        self.regression_head = nn.Linear(10, 1)\n",
    "        self.classification_head = nn.Linear(10, len(np.unique(labels[:, :, 0])))\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = torch.relu(x)\n",
    "        x = self.pool(x)\n",
    "        x = self.flatten(x)\n",
    "        x = self.fc1(x)\n",
    "        x = torch.relu(x)\n",
    "        regression_output = self.regression_head(x)\n",
    "        classification_output = self.classification_head(x)\n",
    "        classification_output = torch.softmax(classification_output, dim=1)\n",
    "        \n",
    "        return regression_output, classification_output\n",
    "\n",
    "model = MyMultiTaskModel()\n",
    "criterion_regression = nn.MSELoss()\n",
    "criterion_classification = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "train_writer = SummaryWriter(log_dir='./logs/train')\n",
    "val_writer = SummaryWriter(log_dir='./logs/validation')\n",
    "\n",
    "\n",
    "def train(model, train_loader, criterion_regression, criterion_classification, optimizer, epoch):\n",
    "    model.train()\n",
    "    running_loss_regression = 0.0\n",
    "    running_loss_classification = 0.0\n",
    "    for i, (inputs, targets_regression, targets_classification) in enumerate(train_loader):\n",
    "        optimizer.zero_grad()\n",
    "        outputs_regression, outputs_classification = model(inputs)\n",
    "        loss_regression = criterion_regression(outputs_regression, targets_regression)\n",
    "        loss_classification = criterion_classification(outputs_classification, targets_classification)\n",
    "        loss = loss_regression + loss_classification\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss_regression += loss_regression.item()\n",
    "        running_loss_classification += loss_classification.item()\n",
    "    avg_loss_regression = running_loss_regression / len(train_loader)\n",
    "    avg_loss_classification = running_loss_classification / len(train_loader)\n",
    "    train_writer.add_scalar('Loss/Regression', avg_loss_regression, epoch)\n",
    "    train_writer.add_scalar('Loss/Classification', avg_loss_classification, epoch)\n",
    "    print(f'Epoch {epoch+1}, Training loss - Regression: {avg_loss_regression}, Classification: {avg_loss_classification}')\n",
    "\n",
    "def validate(model, validation_loader, criterion_regression, criterion_classification, epoch):\n",
    "    model.eval()\n",
    "    validation_loss_regression = 0.0\n",
    "    validation_loss_classification = 0.0\n",
    "    with torch.no_grad():\n",
    "        for inputs, targets_regression, targets_classification in validation_loader:\n",
    "            outputs_regression, outputs_classification = model(inputs)\n",
    "            loss_regression = criterion_regression(outputs_regression, targets_regression)\n",
    "            loss_classification = criterion_classification(outputs_classification, targets_classification)\n",
    "            validation_loss_regression += loss_regression.item()\n",
    "            validation_loss_classification += loss_classification.item()\n",
    "    avg_val_loss_regression = validation_loss_regression / len(validation_loader)\n",
    "    avg_val_loss_classification = validation_loss_classification / len(validation_loader)\n",
    "    val_writer.add_scalar('Loss/Regression', avg_val_loss_regression, epoch)\n",
    "    val_writer.add_scalar('Loss/Classification', avg_val_loss_classification, epoch)\n",
    "    print(f'Validation loss - Regression: {avg_val_loss_regression}, Classification: {avg_val_loss_classification}')\n",
    "\n",
    "def close_writers():\n",
    "    train_writer.close()\n",
    "    val_writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b353c03d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Training loss - Regression: 0.0936830411418582, Classification: 0.6499574985243066\n",
      "Validation loss - Regression: 0.0752740916637829, Classification: 0.5798691761894867\n",
      "Epoch 2, Training loss - Regression: 0.049434231372021145, Classification: 0.5792443142255721\n",
      "Validation loss - Regression: 0.059954661176572714, Classification: 0.5687728356570005\n"
     ]
    }
   ],
   "source": [
    "set_seed(42)\n",
    "num_epochs = 2\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    train(model, train_loader, criterion_regression, criterion_classification, optimizer, epoch)\n",
    "    validate(model, validation_loader, criterion_regression, criterion_classification, epoch)\n",
    "\n",
    "close_writers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "72c29bce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MTL-CNN的回归R平方值为 0.9530041482657609\n",
      "MTL-CNN的回归均方误差为 0.060007494\n",
      "MTL-CNN的回归平均绝对误差为 0.118410505\n",
      "******************************************************\n",
      "MTL-CNN的分类准确率为 0.983453670276775\n",
      "MTL-CNN的分类精确率为 0.9838066626284799\n",
      "MTL-CNN的分类召回率为 0.983453670276775\n",
      "MTL-CNN的分类F1分数为 0.9835551288053177\n",
      "******************************************************\n"
     ]
    }
   ],
   "source": [
    "from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error, accuracy_score, precision_score, recall_score, f1_score\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    y_validation_pred_regression, y_validation_pred_classification = model(validation_features)\n",
    "    y_validation_pred_regression = y_validation_pred_regression.numpy()\n",
    "    y_validation_pred_classification = y_validation_pred_classification.argmax(dim=1).numpy()\n",
    "\n",
    "from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error\n",
    "\n",
    "r2 = r2_score(validation_regression_labels, y_validation_pred_regression)\n",
    "mse = mean_squared_error(validation_regression_labels, y_validation_pred_regression)\n",
    "mae = mean_absolute_error(validation_regression_labels, y_validation_pred_regression)\n",
    "\n",
    "print('MTL-CNN的回归R平方值为', r2)\n",
    "print('MTL-CNN的回归均方误差为', mse)\n",
    "print('MTL-CNN的回归平均绝对误差为', mae)\n",
    "print(\"******************************************************\")\n",
    "\n",
    "accuracy = accuracy_score(validation_classification_labels, y_validation_pred_classification)\n",
    "precision = precision_score(validation_classification_labels, y_validation_pred_classification, average='weighted')\n",
    "recall = recall_score(validation_classification_labels, y_validation_pred_classification, average='weighted')\n",
    "f1 = f1_score(validation_classification_labels, y_validation_pred_classification, average='weighted')\n",
    "\n",
    "print('MTL-CNN的分类准确率为', accuracy)\n",
    "print('MTL-CNN的分类精确率为', precision)\n",
    "print('MTL-CNN的分类召回率为', recall)\n",
    "print('MTL-CNN的分类F1分数为', f1)\n",
    "print(\"******************************************************\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b67f9c4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# 保存模型参数\n",
    "torch.save(model.state_dict(), 'model_weights.pth')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ae9ba4f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import json\n",
    "\n",
    "model.load_state_dict(torch.load('model_weights.pth'))\n",
    "\n",
    "# 提取模型参数\n",
    "model_params = model.state_dict()\n",
    "\n",
    "# 将参数转换为Python的可序列化类型（如列表）\n",
    "params_dict = {key: value.tolist() for key, value in model_params.items()}\n",
    "\n",
    "# 将参数保存为格式化的 JSON 文件\n",
    "with open('model_weights.json', 'w') as f:\n",
    "    json.dump(params_dict, f, indent=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9a079737",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "模型已成功导出为ONNX格式\n"
     ]
    }
   ],
   "source": [
    "dummy_input = torch.randn(1, 8, 10)  # 根据你的模型输入形状进行调整\n",
    "\n",
    "torch.onnx.export(\n",
    "    model,                     # 你的训练过的模型\n",
    "    dummy_input,               # 示例输入\n",
    "    \"model.onnx\",              # 输出文件名\n",
    "    export_params=True,        # 导出模型的所有参数\n",
    "    opset_version=10,          # 指定ONNX的opset版本\n",
    "    do_constant_folding=True,  # 是否执行常量折叠优化\n",
    "    input_names=['input'],     # 输入的名称\n",
    "    output_names=['regression_output', 'classification_output'],  # 输出的名称\n",
    "    dynamic_axes={'input': {0: 'batch_size'},    # 设定动态维度\n",
    "                  'regression_output': {0: 'batch_size'},\n",
    "                  'classification_output': {0: 'batch_size'}}\n",
    ")\n",
    "\n",
    "print(\"模型已成功导出为ONNX格式\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e96f786b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting onnx-tfNote: you may need to restart the kernel to use updated packages.\n",
      "\n",
      "  Using cached onnx_tf-1.10.0-py3-none-any.whl.metadata (510 bytes)\n",
      "Collecting onnx>=1.10.2 (from onnx-tf)\n",
      "  Downloading onnx-1.16.2-cp311-cp311-win_amd64.whl.metadata (16 kB)\n",
      "Collecting PyYAML (from onnx-tf)\n",
      "  Downloading PyYAML-6.0.2-cp311-cp311-win_amd64.whl.metadata (2.1 kB)\n",
      "Collecting tensorflow-addons (from onnx-tf)\n",
      "  Using cached tensorflow_addons-0.22.0-cp311-cp311-win_amd64.whl.metadata (1.8 kB)\n",
      "Requirement already satisfied: numpy>=1.20 in d:\\anaconda\\envs\\tensorflow\\lib\\site-packages (from onnx>=1.10.2->onnx-tf) (1.26.4)\n",
      "Requirement already satisfied: protobuf>=3.20.2 in d:\\anaconda\\envs\\tensorflow\\lib\\site-packages (from onnx>=1.10.2->onnx-tf) (4.25.4)\n",
      "Collecting typeguard<3.0.0,>=2.7 (from tensorflow-addons->onnx-tf)\n",
      "  Using cached typeguard-2.13.3-py3-none-any.whl.metadata (3.6 kB)\n",
      "Requirement already satisfied: packaging in d:\\anaconda\\envs\\tensorflow\\lib\\site-packages (from tensorflow-addons->onnx-tf) (24.1)\n",
      "Using cached onnx_tf-1.10.0-py3-none-any.whl (226 kB)\n",
      "Downloading onnx-1.16.2-cp311-cp311-win_amd64.whl (14.4 MB)\n",
      "   ---------------------------------------- 0.0/14.4 MB ? eta -:--:--\n",
      "   ---------------------------------------- 0.0/14.4 MB ? eta -:--:--\n",
      "   - -------------------------------------- 0.5/14.4 MB 2.8 MB/s eta 0:00:05\n",
      "   ------ --------------------------------- 2.4/14.4 MB 8.4 MB/s eta 0:00:02\n",
      "   ------ --------------------------------- 2.4/14.4 MB 8.4 MB/s eta 0:00:02\n",
      "   ------- -------------------------------- 2.9/14.4 MB 4.2 MB/s eta 0:00:03\n",
      "   -------- ------------------------------- 3.1/14.4 MB 3.8 MB/s eta 0:00:04\n",
      "   --------- ------------------------------ 3.4/14.4 MB 3.1 MB/s eta 0:00:04\n",
      "   ---------- ----------------------------- 3.9/14.4 MB 2.9 MB/s eta 0:00:04\n",
      "   ----------- ---------------------------- 4.2/14.4 MB 2.6 MB/s eta 0:00:04\n",
      "   ------------- -------------------------- 4.7/14.4 MB 2.6 MB/s eta 0:00:04\n",
      "   ------------- -------------------------- 4.7/14.4 MB 2.6 MB/s eta 0:00:04\n",
      "   -------------- ------------------------- 5.2/14.4 MB 2.5 MB/s eta 0:00:04\n",
      "   --------------- ------------------------ 5.5/14.4 MB 2.4 MB/s eta 0:00:04\n",
      "   --------------- ------------------------ 5.8/14.4 MB 2.3 MB/s eta 0:00:04\n",
      "   ---------------- ----------------------- 6.0/14.4 MB 2.2 MB/s eta 0:00:04\n",
      "   ------------------ --------------------- 6.6/14.4 MB 2.1 MB/s eta 0:00:04\n",
      "   ------------------ --------------------- 6.8/14.4 MB 2.1 MB/s eta 0:00:04\n",
      "   ------------------- -------------------- 7.1/14.4 MB 2.1 MB/s eta 0:00:04\n",
      "   -------------------- ------------------- 7.3/14.4 MB 2.0 MB/s eta 0:00:04\n",
      "   -------------------- ------------------- 7.3/14.4 MB 2.0 MB/s eta 0:00:04\n",
      "   --------------------- ------------------ 7.9/14.4 MB 1.9 MB/s eta 0:00:04\n",
      "   ---------------------- ----------------- 8.1/14.4 MB 1.9 MB/s eta 0:00:04\n",
      "   ----------------------- ---------------- 8.4/14.4 MB 1.9 MB/s eta 0:00:04\n",
      "   ----------------------- ---------------- 8.7/14.4 MB 1.9 MB/s eta 0:00:04\n",
      "   ------------------------ --------------- 8.9/14.4 MB 1.9 MB/s eta 0:00:03\n",
      "   ------------------------- -------------- 9.2/14.4 MB 1.8 MB/s eta 0:00:03\n",
      "   -------------------------- ------------- 9.7/14.4 MB 1.8 MB/s eta 0:00:03\n",
      "   --------------------------- ------------ 10.0/14.4 MB 1.8 MB/s eta 0:00:03\n",
      "   ----------------------------- ---------- 10.5/14.4 MB 1.8 MB/s eta 0:00:03\n",
      "   ----------------------------- ---------- 10.7/14.4 MB 1.8 MB/s eta 0:00:03\n",
      "   ------------------------------- -------- 11.3/14.4 MB 1.9 MB/s eta 0:00:02\n",
      "   -------------------------------- ------- 11.8/14.4 MB 1.9 MB/s eta 0:00:02\n",
      "   ---------------------------------- ----- 12.6/14.4 MB 1.9 MB/s eta 0:00:01\n",
      "   ------------------------------------ --- 13.1/14.4 MB 2.0 MB/s eta 0:00:01\n",
      "   ------------------------------------- -- 13.6/14.4 MB 2.0 MB/s eta 0:00:01\n",
      "   -------------------------------------- - 13.9/14.4 MB 2.0 MB/s eta 0:00:01\n",
      "   ---------------------------------------- 14.4/14.4 MB 2.0 MB/s eta 0:00:00\n",
      "Downloading PyYAML-6.0.2-cp311-cp311-win_amd64.whl (161 kB)\n",
      "Using cached tensorflow_addons-0.22.0-cp311-cp311-win_amd64.whl (719 kB)\n",
      "Using cached typeguard-2.13.3-py3-none-any.whl (17 kB)\n",
      "Installing collected packages: typeguard, PyYAML, onnx, tensorflow-addons, onnx-tf\n",
      "Successfully installed PyYAML-6.0.2 onnx-1.16.2 onnx-tf-1.10.0 tensorflow-addons-0.22.0 typeguard-2.13.3\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  WARNING: The scripts backend-test-tools.exe, check-model.exe and check-node.exe are installed in 'D:\\ANACONDA\\envs\\tensorflow\\Scripts' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\n",
      "  WARNING: The script onnx-tf.exe is installed in 'D:\\ANACONDA\\envs\\tensorflow\\Scripts' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\n"
     ]
    }
   ],
   "source": [
    "pip install onnx-tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "066e5a81",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as gen_tensor_dict while saving (showing 1 of 1). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: model_tf\\assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: model_tf\\assets\n"
     ]
    }
   ],
   "source": [
    "from onnx_tf.backend import prepare\n",
    "import onnx\n",
    "\n",
    "onnx_model = onnx.load(\"model.onnx\")\n",
    "tf_rep = prepare(onnx_model) \n",
    "tf_rep.export_graph(\"model_tf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1b4c4b21-e17c-4070-ae90-0e77cb42e0b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:Found untraced functions such as gen_tensor_dict while saving (showing 1 of 1). These functions will not be directly callable after loading.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: model.pb\\assets\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Assets written to: model.pb\\assets\n"
     ]
    }
   ],
   "source": [
    "import onnx\n",
    "from onnx_tf.backend import prepare\n",
    "import tensorflow as tf\n",
    "\n",
    "TF_PATH = \"model.pb\"  # 转tensorflow的pb文件的存储路径\n",
    "ONNX_PATH = \"model.onnx\"  # 需要转换的onnx文件路径\n",
    "TFLITE_PATH = \"model.tflite\"  # 输出的tflite文件路径\n",
    "\n",
    "\"\"\"onnx=>.pb文件\"\"\"\n",
    "onnx_model = onnx.load(ONNX_PATH)  # load onnx model\n",
    "tf_rep = prepare(onnx_model)  # creating TensorflowRep object\n",
    "tf_rep.export_graph(TF_PATH)\n",
    "\n",
    "\"\"\".pb=>tflite\"\"\"\n",
    "converter = tf.lite.TFLiteConverter.from_saved_model(TF_PATH)\n",
    "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
    "tf_lite_model = converter.convert()\n",
    "with open(TFLITE_PATH, 'wb') as f:\n",
    "    f.write(tf_lite_model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d608e90e-42c1-46fd-9e45-f11fcd86d7f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install tensorflow==2.15.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4511f4ef-f7a7-4c62-9f27-68f8eedf1163",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
